import os
import argparse
import json

from PIL import Image

import torch
from torchvision.transforms.functional import to_tensor, to_pil_image

from model import Generator


torch.backends.cudnn.enabled = False
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True


def load_image(image_path, x32=False):
    img = Image.open(image_path).convert("RGB")

    if x32:
        def to_32s(x):
            return 256 if x < 256 else x - x % 32
        w, h = img.size
        img = img.resize((to_32s(w), to_32s(h)))

    return img


def doCreateKaToonImage(args):
    device = args.device
    net = Generator()
    net.load_state_dict(torch.load(args.checkpoint, map_location="cpu"))
    net.to(device).eval()
    os.makedirs(args.output_dir, exist_ok=True)
    imagePath=args.input_dir+'/'+args.image_name
    isHaveImage=args.image_name!='' and os.path.exists(imagePath)
    if isHaveImage:
        if os.path.splitext(args.image_name)[-1].lower() not in [".jpg", ".jpeg",".png", ".bmp", ".tiff",".jfif"]:
            return json.dumps({'code':500,'message':'this image type is not supported','data':''})
        else:
            image=load_image(imagePath, args.x32)
            with torch.no_grad():
                image = to_tensor(image).unsqueeze(0) * 2 - 1
                out = net(image.to(device), args.upsample_align).cpu()
                out = out.squeeze(0).clip(-1, 1) * 0.5 + 0.5
                out = to_pil_image(out)
            out.save(os.path.join(args.output_dir, args.image_name))
            return json.dumps({'code':200,'message':'SUCCESS','data':args.output_dir+"/"+args.image_name})
    else:
        return json.dumps({'code':500,'message':'the file path does not exist','data':''})


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    # 训练集，改版模型得到不同的漫画图片
    parser.add_argument(
        '--checkpoint',
        type=str,
        default='./weights/qingxin.pt',
    )
    # 图片名称
    parser.add_argument(
        '--image_name', 
        type=str, 
        default='',
    )
    # 图片输入文件夹
    parser.add_argument(
        '--input_dir', 
        type=str, 
        default='./samples/inputs',
    )
    # 图片输出文件夹
    parser.add_argument(
        '--output_dir', 
        type=str, 
        default='./samples/results',
    )

    # 是否使用gpu或者cup
    parser.add_argument(
        '--device',
        type=str,
        default='cpu',
    )
    parser.add_argument(
        '--upsample_align',
        type=bool,
        default=False,
        help="Align corners in decoder upsampling layers"
    )
    parser.add_argument(
        '--x32',
        action="store_true",
        help="Resize images to multiple of 32"
    )
    args = parser.parse_args()
    
    print(doCreateKaToonImage(args))
